package com.whosly.trino.client;

import io.trino.client.JsonCodec;
import io.trino.client.JsonResponse;
import io.trino.client.QueryResults;
import okhttp3.OkHttpClient;
import okhttp3.Request;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.OptionalLong;
import java.util.concurrent.TimeUnit;

import static io.trino.client.JsonCodec.jsonCodec;

/**
 * @author fengyang(yueny09 @ 163.com)
 * @date 2022年05月22日 下午4:57
 */
public class TrinoClient implements TClient {
    private static final Logger logger = LoggerFactory.getLogger(TrinoClient.class);

    private static final JsonCodec<QueryResults> QUERY_RESULTS_CODEC = jsonCodec(QueryResults.class);

    private OkHttpClient.Builder clientBuilder;
    private TrinoConfig config;

    public TrinoClient() {
        OkHttpClient.Builder clientBuilder = new OkHttpClient.Builder();
        // 创建客户端，这里设置了 readTimeout 和 callTimeout
        // 设置的目的是 debug Coordinator 的时候客户端会超时失败，导致后续无法进行
        OkHttpClient.Builder client = clientBuilder
                .callTimeout(1, TimeUnit.DAYS)
                .readTimeout(1, TimeUnit.DAYS);

        this.clientBuilder = client;
    }

    public TrinoClient(TrinoConfig config) {
        this();

        this.config = config;
    }

    @Override
    public Request.Builder getRequestBuilder(){
        Request.Builder builder = new Request.Builder()
                .addHeader("X-Trino-User", "trino")
                .addHeader("User-Agent", "StatementClientV1/356");

        return builder;
    }

    @Override
    public JsonResponse<QueryResults> execute(Request request){
        OkHttpClient client = clientBuilder.build();

        // 接受 QueryResult 的返回结果
        JsonResponse<QueryResults> response;
        // 执行请求将结果给 response
        response = JsonResponse.execute(QUERY_RESULTS_CODEC, client, request, OptionalLong.empty());

        logger.trace("TrinoClient execute, req:{}, result:{}.", request, response.getResponseBody());

        return response;
    }

    @Override
    public String getServerAddress() {
        return config != null ? config.getServerAddress() : "";
    }

}
